﻿#pragma once

template <class K, class V>
struct AVLTreeNode {
	AVLTreeNode<K, V>* _left;
	AVLTreeNode<K, V>* _right;
	AVLTreeNode<K, V>* _parent;
	pair<K, V> _kv;
	int _bf;//平衡因子

	AVLTreeNode(const pair<K, V>& kv)
		:_left(nullptr)
		, _right(nullptr)
		, _parent(nullptr)
		, _kv(kv)
		, _bf(0) {
	}
};

template <class K, class V>
struct AVLTree {
	typedef AVLTreeNode<K, V> Node;
public:
	bool Insert(const pair<K, V>& kv) {
		if (_root == nullptr) {
			_root = new Node(kv);
			return true;
		}
		Node* parent = nullptr, * cur = _root;
		while (cur) {
			if (kv.first > cur->_kv.first) {
				parent = cur;
				cur = cur->_right;
			}
			else if (kv.first < cur->_kv.first) {
				parent = cur;
				cur = cur->_left;
			}
			else
				return false;
		}
		cur = new Node(kv);
		if (kv.first > parent->_kv.first) {
			parent->_right = cur;
			cur->_parent = parent;
		}
		else {
			parent->_left = cur;
			cur->_parent = parent;
		}
		//调整平衡
		//新增节点可能会影响祖先，取决于子树的高度是否变化
		//1.如果子树高度[不变]，就[不会]继续往上影响祖先
		//2.如果子树高度[变化]，就[会]继续往上影响祖先
		//新增节点在左子树，父节点_bf--；新增节点在右子树，父节点_bf++。
		//是否要继续往上更新，分3种情况
		//1.父节点_bf更新后 == 0，父节点的子树高度不变，不用继续往上更新，插入结束
		//  ps：在插入之前，父节点_bf == -1 或 1，一边高一边低，新插入节点填上低的那边
		//2.父节点_bf更新后 == 1 或 -1，父节点的子树高度变化，必须继续往上更新
		//  ps：在插入之前，父节点_bf == 0，两边一样高，插入导致高度变化
		//3.父节点_bf更新后 == 2 或 -2，父节点的子树违反规则，需要调整

		//往上查看最高需要到根，所以cur需要指向根，此时parent就会空，所以parent不为空就继续
		while (parent) {
			//更新父节点平衡因子
			if (cur == parent->_left) parent->_bf--;
			else parent->_bf++;
			//查看父节点平衡因子，以此判断是否需要调整
			if (parent->_bf == 0) break;
			else if (parent->_bf == 1 || parent->_bf == -1) {
				cur = parent;
				parent = parent->_parent;
			}
			else if (parent->_bf == 2 || parent->_bf == -2) {
				//旋转
				//1.保证左右均衡
				//2.保证二叉搜索树的规则
				if (parent->_bf == 2 && cur->_bf == 1) //该平衡因子特征是左单旋
					RotateL(parent);
				else if (parent->_bf == -2 && cur->_bf == -1) //该平衡因子特征是右单旋
					RotateR(parent);
				else if (parent->_bf == 2 && cur->_bf == -1)
					RotateRL(parent);
				else if (parent->_bf == -2 && cur->_bf == 1)
					RotateLR(parent);
				//旋转后不需要再往上查看，因为新插入节点增加了高度，但旋转又减低了高度
				break;
			}
			else //以防平衡因子超过2，不是AVLTree，直接报错
				assert(false);
		}
		return true;
	}

	//左单旋
	//解释1：
	//假设平衡因子是2的节点为A，它的右子节点为B
	//将A->right指向B->left，在将B->left指向A
	//解释2：
	//如果一个节点的父节点平衡因子是2，将它的左子结点给其父节点的右子节点
	//再将其父节点当做自己的左子结点

	//左单旋场景，是指新节点插入在较高右子树的右侧，导致右子树的高度比左子树高超过1 
	//左单旋通过将较高的右子树提升为新的根节点，同时调整其左子树的位置，使得左右子树的高度差减小
	void RotateL(Node* parent) {
		//parent是平衡因子为2的节点
		Node* subR = parent->_right;
		Node* subRL = subR->_left;
		Node* parentParent = parent->_parent;//parent不一定是整棵树的根
		//旋转完后，subR是子树的根，还需要处理subR的parent

		//旋转
		parent->_right = subRL;
		subR->_left = parent;
		//上方还未处理完，根节点没有确定，各节点的parent还未更新
		parent->_parent = subR;
		if (subRL) subRL->_parent = parent;

		//处理根节点
		if (_root == parent) { //parent是整棵树的根
			_root = subR;
			subR->_parent = nullptr;
		}
		else { //parent不一定是整棵树的根，旋转完后，还需要处理subR的parent
			if (parent == parentParent->_left) parentParent->_left = subR;
			else parentParent->_right = subR;
			subR->_parent = parentParent;
		}
		parent->_bf = subR->_bf = 0;//处理平衡因子
	}

	void RotateR(Node* parent) {
		Node* subL = parent->_left, * subLR = subL->_right, * parentParent = parent->_parent;

		//旋转
		parent->_left = subLR;
		subL->_right = parent;
		//处理各节点的父节点
		parent->_parent = subL;
		if (subLR) subLR->_parent = parent;

		//处理根节点
		if (_root == parent) {
			_root = subL;
			subL->_parent = nullptr;
		}
		else {
			if (parent == parentParent->_left) parentParent->_left = subL;
			else parentParent->_right = subL;
			subL->_parent = parentParent;
		}
		parent->_bf = subL->_bf = 0;
	}

	//新节点插入较高右子树的左侧---右左：先右单旋再左单旋
	void RotateRL(Node* parent) {
		Node* subR = parent->_right;
		Node* subRL = subR->_left;
		int bf = subRL->_bf;//记录新增节点后subRL节点的平衡因子，用于在旋转后正确更新相关节点的平衡因子。
		//新增节点可能为subRL 或 subRL的左子树 或subRL的右子树。
		//这三种不同的情况会导致旋转完成后，parent、subR、subRL的平衡因子不同

		RotateR(subR);
		RotateL(parent);

		if (bf == 0)
			parent->_bf = subR->_bf = subRL->_bf = 0;//subRL自己是新增
		else if (bf == -1) { //在subRL的左子树新增
			parent->_bf = subRL->_bf = 0;
			subR->_bf = 1;
		}
		else if (bf == 1) { //在subRL的右子树新增
			subR->_bf = subRL->_bf = 0;
			parent->_bf = -1;
		}
		else
			assert(false);
	}

	//新节点插入较高左子树的右侧---左右：先左单旋再右单旋
	void RotateLR(Node* parent) {
		Node* subL = parent->_left;
		Node* subLR = subL->_right;
		int bf = subLR->_bf;

		RotateL(subL);
		RotateR(parent);

		if (bf == 0)
			parent->_bf = subL->_bf = subLR->_bf = 0;
		else if (bf == -1) {
			subL->_bf = subLR->_bf = 0;
			parent->_bf = 1;
		}
		else if (bf == 1) {
			parent->_bf = subLR->_bf = 0;
			subL->_bf = -1;
		}
		else
			assert(false);
	}

	Node* Find(const K& key) {
		Node* cur = _root;
		while (cur) {
			if (key > _root->_kv.first) cur = cur->_right;
			else if (key < _root->_kv.first) cur = cur->_left;
			else return cur;
		}
		return nullptr;
	}

	void InOrder() {
		_InOrder(_root);
		cout << endl;
	}

	bool IsBalance() { return _IsBalance(_root); }
	int Height() { return _Height(_root); }
	size_t Size() { return _Size(_root); }
private:
	Node* _root = nullptr;

	void _InOrder(Node* root) {
		if (root == nullptr)
			return;
		_InOrder(root->_left);
		cout << root->_kv.first << " ";
		_InOrder(root->_right);
	}

	int _Height(Node* root) {
		if (root == nullptr) return 0;
		int leftHeight = _Height(root->_left);
		int rightHeight = _Height(root->_right);

		return leftHeight > rightHeight ? leftHeight + 1 : rightHeight + 1;
	}
	bool _IsBalance(Node* root) {
		if (root == nullptr) return true;

		int leftHeight = _Height(root->_left);
		int rightHeight = _Height(root->_right);
		//验证当前节点的平衡因子是否正确，通过实际计算得到的左右子树高度差去验证
		//如果不相等，说明平衡因子计算或维护有误
		if (rightHeight - leftHeight != root->_bf) {
			cout << root->_kv.first << "平衡因子异常" << endl;
			return false;
		}

		return abs(rightHeight - leftHeight) < 2 //当前节点左右子树高度
			&& _IsBalance(root->_left) //继续检查其左子树的左右子树高度
			&& _IsBalance(root->_right);
	}

	size_t _Size(Node* root) {
		if (root == nullptr) return 0;
		return _Size(root->_left) + _Size(root->_right) + 1;
	}
};